import torch
import torchvision.models as models
import onnx

model   = models.resnet18(pretrained=True)
# model = models.resnet50(pretrained=True)                   # (1,3,224,224)
# model = models.resnet152(pretrained=True)                # (1,3,224,224)
# model = models.mobilenet.mobilenet_v2(pretrained=True)   # (1,3,112,112)
# model = models.squeezenet.squeezenet1_0(pretrained=True) # (1,3,224,224)
# model = models.shufflenet_v2_x0_5(pretrained=True)       # (1,3,224,224)

dummy  = torch.randn(1, 3, 224, 224)
# dummy  = torch.randn(1,3,112,112)
torch.onnx.export(model, dummy, "resnet18.onnx", input_names=["input"], output_names=["output"])